package org.example.springaitest.retriever;

import org.neo4j.driver.*;
import org.neo4j.driver.Record;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;

import java.util.*;
import java.util.stream.Collectors;

public class Neo4jVectorDocumentRetriever implements DocumentRetriever {

    private final VectorStore vectorStore;
    private final Driver driver;
    private final int topK;

    public Neo4jVectorDocumentRetriever(VectorStore vectorStore, Driver driver, int topK) {
        this.vectorStore = vectorStore;
        this.driver = driver;
        this.topK = topK;
    }

    @Override
    public List<Document> retrieve(Query query) {
        // 1. 向量召回
        List<Document> vectorDocs = vectorStore.similaritySearch(
                SearchRequest.builder()
                        .query(query.text())
                        .topK(topK)
                        .build());

        // 2. 对每个文档做图谱扩展（只改 text）
        return vectorDocs.stream()
                .map(this::expandWithGraph)
                .collect(Collectors.toList());
    }

    private Document expandWithGraph(Document doc) {
        StringBuilder enriched = new StringBuilder();

        try (Session session = driver.session()) {

            // Step 1: 获取 Answer 节点的 text 内容（真实回答）
            Result answerResult = session.run("""
            MATCH (d:Document {id: $id})-[:HAS_ANSWER]->(a:Answer)
            RETURN a.text AS text
            """, Map.of("id", doc.getId()));

            if (!answerResult.hasNext()) {
                // 降级：使用原始 Document 的文本
                enriched.append(doc.getText());
                enriched.append("\n\n[Warning: No Answer node found for Document ID: ").append(doc.getId()).append("]\n");
            } else {
                Record record = answerResult.next();
                String answerText = record.get("text").asString();
                enriched.append(answerText); // ✅ 使用 Answer 的真实回答
            }

            // Step 2: 通过 MENTIONS_ENTITY 找到所有提及的实体（用于图谱扩展）
            Result entityResult = session.run("""
            MATCH (d:Document {id: $id})-[:HAS_ANSWER]->(:Answer)-[:MENTIONS_ENTITY]->(e)
            RETURN e.name AS name, labels(e)[0] AS label
            """, Map.of("id", doc.getId()));

            Set<String> entityNames = new HashSet<>();
            List<String> entityLines = new ArrayList<>();

            while (entityResult.hasNext()) {
                Record r = entityResult.next();
                String name = r.get("name").asString();
                String label = r.get("label").asString();
                entityLines.add(label + ": " + name);
                entityNames.add(name);
            }

            if (!entityNames.isEmpty()) {
                enriched.append("\n\n--- Graph Context ---\n");
                enriched.append("Entities mentioned:\n").append(String.join("\n", entityLines)).append("\n");

                // Step 3: 查询这些实体之间的“真实业务关系”（排除 MENTIONS_ENTITY）
// 把 expandWithGraph 里 Step3 的 Cypher 整块换掉
                Result relResult = session.run("""
    // 1. 先把所有可能实体一次性捞出来（避免 UNWIND 丢数据）
    MATCH (d:Document {id: $id})-[:HAS_ANSWER]->(:Answer)-[:MENTIONS_ENTITY]->(e)
    WITH collect(DISTINCT e) AS entities
    UNWIND entities AS e
    
    // 2. 双向扫关系，不要求其他节点有 name（只要求它有 label 即可）
    MATCH (e)-[r]-(other)
    WHERE NOT type(r) IN ['HAS_ANSWER','MENTIONS_ENTITY']   // 过滤掉边
      AND e <> other                                        // 去掉自环
    RETURN e.name      AS start,
           type(r)     AS rel,
           other.name  AS end,
           labels(e)[0]     AS startLabel,
           labels(other)[0] AS endLabel
    """, Map.of("id", doc.getId()));

                Set<String> printedTriples = new HashSet<>();
                List<String> relationships = new ArrayList<>();

                while (relResult.hasNext()) {
                    Record r = relResult.next();
                    String start = r.get("start").asString();
                    String rel = r.get("rel").asString();
                    String end = r.get("end").asString();

                    String key = start + "-" + rel + "->" + end;
                    if (!printedTriples.contains(key)) {
                        relationships.add(start + " -[" + rel + "]-> " + end);
                        printedTriples.add(key);
                    }
                }

                if (!relationships.isEmpty()) {
                    enriched.append("\nRelated Relationships:\n").append(String.join("\n", relationships)).append("\n");
                } else {
                    enriched.append("\nNo additional business relationships found.\n");
                }
            }
            // Step 4: 获取 Answer 使用的 Table 及其 Columns
            Result tableColumnResult = session.run("""
    MATCH (d:Document {id: $id})-[:HAS_ANSWER]->(a:Answer)-[:USES_TABLE]->(t:Table)
    OPTIONAL MATCH (t)-[:HAS_COLUMN]->(c:Column)
    RETURN t.name AS tableName,
           collect({
               name: c.name,
               comment: c.comment
           }) AS columns
    """, Map.of("id", doc.getId()));

            List<String> tableColumnLines = new ArrayList<>();
            while (tableColumnResult.hasNext()) {
                Record r = tableColumnResult.next();
                String tableName = r.get("tableName").asString();

                // ✅ 正确转换：将每个元素作为 map 解析
                List<Map<String, Object>> columns = r.get("columns").asList(Value::asMap);

                List<String> columnDescriptions = new ArrayList<>();
                for (Map<String, Object> col : columns) {
                    String colName = Optional.ofNullable(col.get("name"))
                            .map(Object::toString)
                            .orElse("<<Unnamed>>");

                    String comment = Optional.ofNullable(col.get("comment"))
                            .map(Object::toString)
                            .map(c -> " COMMENT '" + c + "'")
                            .orElse("");

                    columnDescriptions.add(colName + comment);
                }

                String columnList = String.join(", ", columnDescriptions);
                tableColumnLines.add("Table: " + tableName + " => Columns: [" + columnList + "]");
            }

            if (!tableColumnLines.isEmpty()) {
                enriched.append("\n--- Table & Column Context ---\n");
                enriched.append(String.join("\n", tableColumnLines)).append("\n");
            }

        } catch (Exception e) {
            // 静默失败，不中断主流程
            enriched.append("\n[Graph enrichment failed]\n");
        }
        System.out.println( enriched);
        // ✅ 严格按照你的要求返回：只改 text，metadata 原样保留
        return Document.builder()
                .id(doc.getId())
                .text(enriched.toString())
                .metadata(doc.getMetadata())  // 不修改 metadata
                .build();
    }
}